{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mindnlp PEFT with DNA Language Models\n",
    "This notebook demonstrates how to utilize parameter-efficient fine-tuning techniques (PEFT) from the PEFT library to fine-tune a DNA Language Model (DNA-LM). The fine-tuned DNA-LM will be applied to solve a task from the nucleotide benchmark dataset. Parameter-efficient fine-tuning (PEFT) techniques are crucial for adapting large pre-trained models to specific tasks with limited computational resources."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When using modelarts env, need to remove `RANK_TABLE_FILE` to use `MindSpore` dynamic graph mode.\n",
    "```python\n",
    "import os\n",
    "if \"RANK_TABLE_FILE\" in os.environ:\n",
    "    del os.environ[\"RANK_TABLE_FILE\"]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Import relevant libraries\n",
    "We'll start by importing the required libraries, including the mindnlp library and other dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/numpy/core/getlimits.py:499: UserWarning: The value of the smallest subnormal for <class 'numpy.float64'> type is zero.\n",
      "  setattr(self, word, getattr(machar, word).flat[0])\n",
      "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float64'> type is zero.\n",
      "  return self._float_to_str(self.smallest_subnormal)\n",
      "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/numpy/core/getlimits.py:499: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.\n",
      "  setattr(self, word, getattr(machar, word).flat[0])\n",
      "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.\n",
      "  return self._float_to_str(self.smallest_subnormal)\n",
      "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Building prefix dict from the default dictionary ...\n",
      "Loading model from cache /tmp/jieba.cache\n",
      "Loading model cost 1.419 seconds.\n",
      "Prefix dict has been built successfully.\n"
     ]
    }
   ],
   "source": [
    "import mindspore\n",
    "import mindnlp\n",
    "import numpy as np\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Load models\n",
    "We'll load a pre-trained DNA Language Model, \"SpeciesLM\", that serves as the base for fine-tuning. This is done using the transformers library from HuggingFace.\n",
    "\n",
    "The tokenizer and the model comes from the paper, \"Species-aware DNA language models capture regulatory elements and their evolution\". Paper Link, Code Link. They introduce a species-aware DNA language model, which is trained on more than 800 species spanning over 500 million years of evolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindnlp.transformers import AutoTokenizer, AutoModelForMaskedLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/mindnlp/transformers/tokenization_utils_base.py:1526: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted, and will be then set to `False` by default. \n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MS_ALLOC_CONF]Runtime config:  enable_vmm:True  vmm_align_size:2MB\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"gagneurlab/SpeciesLM\", revision = \"downstream_species_lm\")\n",
    "lm = AutoModelForMaskedLM.from_pretrained(\"gagneurlab/SpeciesLM\", revision = \"downstream_species_lm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BertForMaskedLM(\n",
       "  (bert): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(5504, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0-11): 12 x BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear (768 -> 768)\n",
       "              (key): Linear (768 -> 768)\n",
       "              (value): Linear (768 -> 768)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear (768 -> 768)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear (768 -> 3072)\n",
       "            (intermediate_act_fn): GELU(approximate='none')\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear (3072 -> 768)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (cls): BertOnlyMLMHead(\n",
       "    (predictions): BertLMPredictionHead(\n",
       "      (transform): BertPredictionHeadTransform(\n",
       "        (dense): Linear (768 -> 768)\n",
       "        (transform_act_fn): GELU(approximate='none')\n",
       "        (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      )\n",
       "      (decoder): Linear (768 -> 5504)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Prepare datasets\n",
    "We'll load the `nucleotide_transformer_downstream_tasks` dataset, which contains 18 downstream tasks from the Nucleotide Transformer paper. This dataset provides a consistent genomics benchmark with binary classification tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'train': <mindspore.dataset.engine.datasets_user_defined.GeneratorDataset at 0xfffe3ef8ac40>,\n",
       " 'test': <mindspore.dataset.engine.datasets_user_defined.GeneratorDataset at 0xfffe3ef741f0>}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from mindnlp.dataset import load_dataset\n",
    "raw_data = load_dataset(\"InstaDeepAI/nucleotide_transformer_downstream_tasks\", \"H3\")\n",
    "raw_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use the \"H3\" subset of this dataset, which contains a total of 13,468 rows in the training data, and 1497 rows in the test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train dataset info:\n",
      "dataset column: ['sequence', 'name', 'label']\n",
      "dataset size: 13468\n",
      "dataset batch size: 1\n",
      "\n",
      "test dataset info:\n",
      "dataset column: ['sequence', 'name', 'label']\n",
      "dataset size: 1497\n",
      "dataset batch size: 1\n",
      "\n"
     ]
    }
   ],
   "source": [
    "train_dataset = raw_data['train']\n",
    "test_dataset = raw_data['test']\n",
    "\n",
    "def show_dataset_info(dataset):\n",
    "    print(\"dataset column: {}\".format(dataset.get_col_names()))\n",
    "    print(\"dataset size: {}\".format(dataset.get_dataset_size()))\n",
    "    print(\"dataset batch size: {}\\n\".format(dataset.get_batch_size()))\n",
    "\n",
    "print(\"train dataset info:\")\n",
    "show_dataset_info(train_dataset)\n",
    "print(\"test dataset info:\")\n",
    "show_dataset_info(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset consists of three columns, `sequence`, `name` and `label`. An row in this dataset looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Tensor(shape=[], dtype=String, value= 'TCACTTCGATTATTGAGGCAGTCTTCATTAAAGTTTATTACAATGGATATGGTATCACCAGTCTTGAACCTACAATCATCTATTTTAGGTGAGCTCGTAGGCATTATTGGAAAAGTGTTCTTTCTCTTAATAGAAGAGATTAAATACCCGATAATCACACCCAAAATTATTGTGGATGCCCAGATATCTTCTTGGTCATTGTTTTTTTTCGCTTCAATCTGTAATCTCTCTGCAAAATTTCGGGAGCCAATAGTGACAACATCGTCAATAATAAGTTTGATGGAATCGGAAAAAGATCTTAAAAATGTAAATGAGTATTTCCAAATAATGGCCAAAATGCTCTTTATATTGGAAAATAAAATAGTTGTTTCGCTCTTCGTAGTATTTAACATTTCCGTTCTTATCATTGTAAAGTCTGAGCCATATTCATATGGAAAAGTGCTTTTTAAACCTAGTTCCTCCATATTTTAGTTTTTTATCGATATTGGAAAAAAAAGAGC'), Tensor(shape=[], dtype=String, value= 'YBR063C_YBR063C_367930|0'), Tensor(shape=[], dtype=Int64, value= 0)]\n"
     ]
    }
   ],
   "source": [
    "for data in train_dataset:\n",
    "    print(data)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We split out dataset into training, test, and validation sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train dataset info:\n",
      "dataset column: ['sequence', 'name', 'label']\n",
      "dataset size: 11448\n",
      "dataset batch size: 1\n",
      "\n",
      "valid dataset info:\n",
      "dataset column: ['sequence', 'name', 'label']\n",
      "dataset size: 2020\n",
      "dataset batch size: 1\n",
      "\n"
     ]
    }
   ],
   "source": [
    "train_dataset, valid_dataset = train_dataset.split([0.85, 0.15])\n",
    "\n",
    "print(\"train dataset info:\")\n",
    "show_dataset_info(train_dataset)\n",
    "print(\"valid dataset info:\")\n",
    "show_dataset_info(valid_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we use the tokenizer and a utility function we created, `get_kmers` to generate the final data and labels. The `get_kmers` function is essential for generating overlapping 6-mers needed by the language model (LM). By using k=6 and stride=1, we ensure that the model receives continuous and overlapping subsequences, capturing the local context within the biological sequence for more effective analysis and prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_kmers(seq, k=6, stride=1):\n",
    "    return [seq[i:i + k] for i in range(0, len(seq), stride) if i + k <= len(seq)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "BATCH_SIZE = 16\n",
    "DATASET_LIMIT = 200 # NOTE: This dataset limit is set to 200, so that the training runs faster. It can be set to None to use the\n",
    "                    # entire dataset\n",
    "\n",
    "def process_sequence(data):\n",
    "    sequence = data.tolist()\n",
    "    sequence = \"candida_glabrata \" + \" \".join(get_kmers(sequence))\n",
    "    sequence = tokenizer(sequence)[\"input_ids\"]\n",
    "    return sequence\n",
    "\n",
    "def nucleotide_dataset_process(dataset):\n",
    "    # remove name column\n",
    "    dataset = dataset.project(columns=['sequence', 'label'])\n",
    "\n",
    "    # process sequence\"\"\"  \"\"\"\n",
    "    dataset = dataset.map(process_sequence, input_columns=['sequence'], output_columns=['sequence'])\n",
    "\n",
    "    # change dataset size\n",
    "    dataset = dataset.take(DATASET_LIMIT)\n",
    "\n",
    "    # batch with padding\n",
    "    dataset = dataset.padded_batch(batch_size=BATCH_SIZE,\n",
    "                                   drop_remainder=True,\n",
    "                                   pad_info={'sequence':(None, -100)})\n",
    "    return dataset\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we create a Dataset object for each our sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train dataset info:\n",
      "dataset column: ['sequence', 'label']\n",
      "dataset size: 12\n",
      "dataset batch size: 16\n",
      "\n",
      "valid dataset info:\n",
      "dataset column: ['sequence', 'label']\n",
      "dataset size: 12\n",
      "dataset batch size: 16\n",
      "\n",
      "test dataset info:\n",
      "dataset column: ['sequence', 'label']\n",
      "dataset size: 12\n",
      "dataset batch size: 16\n",
      "\n",
      "[Tensor(shape=[16, 498], dtype=Int64, value=\n",
      "[[   2, 4724,  200 ... 3930, 3420,    3],\n",
      " [   2, 4724,  414 ...  518, 2057,    3],\n",
      " [   2, 4724,    8 ... 3674, 2394,    3],\n",
      " ...\n",
      " [   2, 4724, 1369 ... 3089,   54,    3],\n",
      " [   2, 4724, 3365 ... 2606, 2219,    3],\n",
      " [   2, 4724, 1630 ... 2727, 2702,    3]]), Tensor(shape=[16], dtype=Int64, value= [1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 1, 1])]\n"
     ]
    }
   ],
   "source": [
    "train_dataset = train_dataset.apply(nucleotide_dataset_process)\n",
    "valid_dataset = valid_dataset.apply(nucleotide_dataset_process)\n",
    "test_dataset = test_dataset.apply(nucleotide_dataset_process)\n",
    "\n",
    "print(\"train dataset info:\")\n",
    "show_dataset_info(train_dataset)\n",
    "\n",
    "print(\"valid dataset info:\")\n",
    "show_dataset_info(valid_dataset)\n",
    "\n",
    "print(\"test dataset info:\")\n",
    "show_dataset_info(test_dataset)\n",
    "\n",
    "for data in train_dataset:\n",
    "    print(data)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Train model\n",
    "Now, we'll train our DNA Language Model with the training dataset. We'll add a linear layer in the final layer of our language model, and then, train all the parameteres of our model with the training dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DNA_LM(\n",
       "  (model): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(5504, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0-11): 12 x BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear (768 -> 768)\n",
       "              (key): Linear (768 -> 768)\n",
       "              (value): Linear (768 -> 768)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear (768 -> 768)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear (768 -> 3072)\n",
       "            (intermediate_act_fn): GELU(approximate='none')\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear (3072 -> 768)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (classifier): Linear (768 -> 2)\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import mindspore\n",
    "from mindnlp.core import nn\n",
    "\n",
    "class DNA_LM(nn.Module):\n",
    "    def __init__(self, model, num_labels):\n",
    "        super(DNA_LM, self).__init__()\n",
    "        self.model = model.bert\n",
    "        self.in_features = model.config.hidden_size\n",
    "        self.out_features = num_labels\n",
    "        self.classifier = nn.Linear(self.in_features, self.out_features)\n",
    "\n",
    "    def forward(self, sequence, label=None):\n",
    "        outputs = self.model(input_ids=sequence, attention_mask=None, output_hidden_states=True)\n",
    "        sequence_output = outputs.hidden_states[-1]\n",
    "        # Use the [CLS] token for classification\n",
    "        cls_output = sequence_output[:, 0, :]\n",
    "        logits = self.classifier(cls_output)\n",
    "\n",
    "        loss = None\n",
    "        if label is not None:\n",
    "            loss_fct = nn.CrossEntropyLoss()\n",
    "            loss = loss_fct(logits.view(-1, self.out_features), label.view(-1))\n",
    "\n",
    "        return (loss, logits) if loss is not None else logits\n",
    "\n",
    "# Number of classes for your classification task\n",
    "num_labels = 2\n",
    "classification_model = DNA_LM(lm, num_labels)\n",
    "classification_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/60 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 1/60 [00:22<21:42, 22.08s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 1/60 [01:06<21:42, 22.08s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8252, 'learning_rate': 1.9666666666666666e-05, 'epoch': 0.08}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|▎         | 2/60 [01:07<34:40, 35.87s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7127, 'learning_rate': 1.9333333333333333e-05, 'epoch': 0.17}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▌         | 3/60 [01:08<18:47, 19.78s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7957, 'learning_rate': 1.9e-05, 'epoch': 0.25}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  7%|▋         | 4/60 [01:09<11:25, 12.24s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7914, 'learning_rate': 1.866666666666667e-05, 'epoch': 0.33}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  8%|▊         | 5/60 [01:09<07:22,  8.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.9366, 'learning_rate': 1.8333333333333333e-05, 'epoch': 0.42}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|█         | 6/60 [01:10<04:58,  5.53s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5876, 'learning_rate': 1.8e-05, 'epoch': 0.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 12%|█▏        | 7/60 [01:10<03:28,  3.93s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5856, 'learning_rate': 1.7666666666666668e-05, 'epoch': 0.58}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 13%|█▎        | 8/60 [01:11<02:29,  2.88s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6507, 'learning_rate': 1.7333333333333336e-05, 'epoch': 0.67}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 15%|█▌        | 9/60 [01:12<01:51,  2.19s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8005, 'learning_rate': 1.7e-05, 'epoch': 0.75}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 17%|█▋        | 10/60 [01:12<01:25,  1.71s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7325, 'learning_rate': 1.6666666666666667e-05, 'epoch': 0.83}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 18%|█▊        | 11/60 [01:13<01:07,  1.38s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6998, 'learning_rate': 1.6333333333333335e-05, 'epoch': 0.92}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 12/60 [01:14<00:55,  1.16s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8102, 'learning_rate': 1.6000000000000003e-05, 'epoch': 1.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      " 20%|██        | 12/60 [01:18<00:55,  1.16s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.6935629844665527, 'eval_runtime': 4.0479, 'eval_samples_per_second': 2.964, 'eval_steps_per_second': 0.494, 'epoch': 1.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 22%|██▏       | 13/60 [01:19<01:47,  2.29s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6874, 'learning_rate': 1.5666666666666667e-05, 'epoch': 1.08}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 23%|██▎       | 14/60 [01:19<01:23,  1.81s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5356, 'learning_rate': 1.5333333333333334e-05, 'epoch': 1.17}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 25%|██▌       | 15/60 [01:20<01:07,  1.50s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.717, 'learning_rate': 1.5000000000000002e-05, 'epoch': 1.25}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 27%|██▋       | 16/60 [01:21<00:57,  1.30s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.711, 'learning_rate': 1.4666666666666666e-05, 'epoch': 1.33}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 28%|██▊       | 17/60 [01:22<00:50,  1.17s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7494, 'learning_rate': 1.4333333333333334e-05, 'epoch': 1.42}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███       | 18/60 [01:23<00:44,  1.07s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5204, 'learning_rate': 1.4e-05, 'epoch': 1.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 32%|███▏      | 19/60 [01:23<00:39,  1.03it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5145, 'learning_rate': 1.3666666666666667e-05, 'epoch': 1.58}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 33%|███▎      | 20/60 [01:24<00:37,  1.07it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5374, 'learning_rate': 1.3333333333333333e-05, 'epoch': 1.67}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 35%|███▌      | 21/60 [01:25<00:34,  1.13it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7067, 'learning_rate': 1.3000000000000001e-05, 'epoch': 1.75}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 37%|███▋      | 22/60 [01:26<00:31,  1.20it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6515, 'learning_rate': 1.2666666666666667e-05, 'epoch': 1.83}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 38%|███▊      | 23/60 [01:26<00:28,  1.30it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6086, 'learning_rate': 1.2333333333333334e-05, 'epoch': 1.92}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 24/60 [01:27<00:28,  1.27it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6736, 'learning_rate': 1.2e-05, 'epoch': 2.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                               \n",
      " 40%|████      | 24/60 [01:45<00:28,  1.27it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.6261480450630188, 'eval_runtime': 4.1712, 'eval_samples_per_second': 2.877, 'eval_steps_per_second': 0.479, 'epoch': 2.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 42%|████▏     | 25/60 [01:46<03:34,  6.12s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5717, 'learning_rate': 1.1666666666666668e-05, 'epoch': 2.08}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 26/60 [01:46<02:32,  4.49s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5132, 'learning_rate': 1.1333333333333334e-05, 'epoch': 2.17}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 45%|████▌     | 27/60 [01:47<01:50,  3.35s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6381, 'learning_rate': 1.1000000000000001e-05, 'epoch': 2.25}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 47%|████▋     | 28/60 [01:48<01:21,  2.53s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.576, 'learning_rate': 1.0666666666666667e-05, 'epoch': 2.33}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 48%|████▊     | 29/60 [01:48<01:01,  1.98s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6151, 'learning_rate': 1.0333333333333335e-05, 'epoch': 2.42}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|█████     | 30/60 [01:49<00:47,  1.59s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4419, 'learning_rate': 1e-05, 'epoch': 2.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 52%|█████▏    | 31/60 [01:50<00:39,  1.38s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4393, 'learning_rate': 9.666666666666667e-06, 'epoch': 2.58}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 53%|█████▎    | 32/60 [01:51<00:33,  1.20s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4634, 'learning_rate': 9.333333333333334e-06, 'epoch': 2.67}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 55%|█████▌    | 33/60 [01:51<00:28,  1.07s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.612, 'learning_rate': 9e-06, 'epoch': 2.75}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 57%|█████▋    | 34/60 [01:52<00:25,  1.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5955, 'learning_rate': 8.666666666666668e-06, 'epoch': 2.83}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 58%|█████▊    | 35/60 [01:53<00:23,  1.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4584, 'learning_rate': 8.333333333333334e-06, 'epoch': 2.92}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 36/60 [01:54<00:20,  1.16it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6288, 'learning_rate': 8.000000000000001e-06, 'epoch': 3.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                               \n",
      " 60%|██████    | 36/60 [02:12<00:20,  1.16it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.5678010582923889, 'eval_runtime': 4.0958, 'eval_samples_per_second': 2.93, 'eval_steps_per_second': 0.488, 'epoch': 3.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 62%|██████▏   | 37/60 [02:13<02:27,  6.42s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4614, 'learning_rate': 7.666666666666667e-06, 'epoch': 3.08}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 63%|██████▎   | 38/60 [02:14<01:43,  4.71s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4509, 'learning_rate': 7.333333333333333e-06, 'epoch': 3.17}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 65%|██████▌   | 39/60 [02:15<01:13,  3.50s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5712, 'learning_rate': 7e-06, 'epoch': 3.25}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 67%|██████▋   | 40/60 [02:15<00:53,  2.69s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5515, 'learning_rate': 6.666666666666667e-06, 'epoch': 3.33}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 68%|██████▊   | 41/60 [02:16<00:40,  2.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6143, 'learning_rate': 6.333333333333333e-06, 'epoch': 3.42}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 70%|███████   | 42/60 [02:17<00:31,  1.73s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.3716, 'learning_rate': 6e-06, 'epoch': 3.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 72%|███████▏  | 43/60 [02:18<00:24,  1.42s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.3609, 'learning_rate': 5.666666666666667e-06, 'epoch': 3.58}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 73%|███████▎  | 44/60 [02:18<00:19,  1.21s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4412, 'learning_rate': 5.333333333333334e-06, 'epoch': 3.67}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 75%|███████▌  | 45/60 [02:19<00:15,  1.05s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6236, 'learning_rate': 5e-06, 'epoch': 3.75}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 77%|███████▋  | 46/60 [02:20<00:13,  1.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5603, 'learning_rate': 4.666666666666667e-06, 'epoch': 3.83}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 78%|███████▊  | 47/60 [02:21<00:11,  1.11it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4473, 'learning_rate': 4.333333333333334e-06, 'epoch': 3.92}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 48/60 [02:21<00:10,  1.15it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5042, 'learning_rate': 4.000000000000001e-06, 'epoch': 4.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                               \n",
      " 80%|████████  | 48/60 [02:39<00:10,  1.15it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.5305204391479492, 'eval_runtime': 4.2506, 'eval_samples_per_second': 2.823, 'eval_steps_per_second': 0.471, 'epoch': 4.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 82%|████████▏ | 49/60 [02:40<01:07,  6.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.385, 'learning_rate': 3.6666666666666666e-06, 'epoch': 4.08}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 83%|████████▎ | 50/60 [02:40<00:44,  4.46s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4341, 'learning_rate': 3.3333333333333333e-06, 'epoch': 4.17}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 85%|████████▌ | 51/60 [02:41<00:29,  3.32s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4644, 'learning_rate': 3e-06, 'epoch': 4.25}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 87%|████████▋ | 52/60 [02:42<00:20,  2.55s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5216, 'learning_rate': 2.666666666666667e-06, 'epoch': 4.33}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 88%|████████▊ | 53/60 [02:42<00:13,  1.98s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5756, 'learning_rate': 2.3333333333333336e-06, 'epoch': 4.42}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 90%|█████████ | 54/60 [02:43<00:09,  1.58s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.3522, 'learning_rate': 2.0000000000000003e-06, 'epoch': 4.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 92%|█████████▏| 55/60 [02:44<00:06,  1.29s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.2875, 'learning_rate': 1.6666666666666667e-06, 'epoch': 4.58}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 93%|█████████▎| 56/60 [02:44<00:04,  1.10s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4042, 'learning_rate': 1.3333333333333334e-06, 'epoch': 4.67}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 95%|█████████▌| 57/60 [02:45<00:02,  1.01it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5759, 'learning_rate': 1.0000000000000002e-06, 'epoch': 4.75}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 97%|█████████▋| 58/60 [02:46<00:01,  1.14it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5405, 'learning_rate': 6.666666666666667e-07, 'epoch': 4.83}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 98%|█████████▊| 59/60 [02:46<00:00,  1.20it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.4023, 'learning_rate': 3.3333333333333335e-07, 'epoch': 4.92}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [02:47<00:00,  1.30it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5196, 'learning_rate': 0.0, 'epoch': 5.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                               \n",
      "100%|██████████| 60/60 [02:51<00:00,  2.86s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.5207817554473877, 'eval_runtime': 4.0536, 'eval_samples_per_second': 2.96, 'eval_steps_per_second': 0.493, 'epoch': 5.0}\n",
      "{'train_runtime': 171.6193, 'train_samples_per_second': 5.594, 'train_steps_per_second': 0.35, 'train_loss': 0.5752782459060352, 'epoch': 5.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=60, training_loss=0.5752782459060352, metrics={'train_runtime': 171.6193, 'train_samples_per_second': 5.594, 'train_steps_per_second': 0.35, 'train_loss': 0.5752782459060352, 'epoch': 5.0})"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from mindnlp.engine import Trainer, TrainingArguments\n",
    "\n",
    "# Define training arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir='./results',\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    num_train_epochs=5,\n",
    "    weight_decay=0.01,\n",
    "    eval_steps=1,\n",
    "    logging_steps=1,\n",
    ")\n",
    "\n",
    "# Initialize Trainer\n",
    "trainer = Trainer(\n",
    "    model=classification_model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=valid_dataset,\n",
    "    tokenizer=tokenizer,\n",
    ")\n",
    "\n",
    "# Train the model\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 12/12 [00:02<00:00,  4.35it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1 1 1 1 0 1 1 0 0 0 0 1 1 0 1 0 1 1 1 0 1 1 0 0 1 0 1 0 0 0 1 1 1 0 0 0 1\n",
      " 1 1 0 0 0 1 1 0 1 0 0 0 0 1 1 1 1 0 0 1 1 0 0 1 1 1 1 0 0 1 0 1 1 0 0 0 1\n",
      " 0 0 0 1 1 0 1 0 0 1 0 0 1 1 1 1 0 0 1 1 1 1 0 1 0 1 1 1 1 0 0 1 1 0 0 0 1\n",
      " 1 0 1 1 0 0 1 1 1 0 1 0 1 1 1 0 0 0 1 0 1 1 1 1 1 1 0 1 1 0 1 0 1 1 0 1 1\n",
      " 0 1 1 1 0 0 1 0 1 1 1 1 1 1 1 1 1 1 0 1 1 1 0 1 0 0 1 1 1 1 1 1 0 0 1 0 0\n",
      " 1 1 0 1 1 0 1]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Generate predictions\n",
    "\n",
    "predictions = trainer.predict(test_dataset)\n",
    "logits = predictions.predictions\n",
    "predicted_labels = logits.argmax(axis=-1)\n",
    "print(predicted_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we create a function to calculate the accuracy from the test and predicted labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.78\n"
     ]
    }
   ],
   "source": [
    "def calculate_accuracy(true_labels, predicted_labels):\n",
    "\n",
    "    assert len(true_labels) == len(predicted_labels), \"Arrays must have the same length\"\n",
    "    correct_predictions = np.sum(true_labels == predicted_labels)\n",
    "    accuracy = correct_predictions / len(true_labels)\n",
    "\n",
    "    return accuracy\n",
    "\n",
    "test_labels = []\n",
    "for data in test_dataset:\n",
    "    sequence, label = data\n",
    "    for single_label in label:\n",
    "        test_labels.append(single_label)\n",
    "\n",
    "accuracy = calculate_accuracy(test_labels, predicted_labels)\n",
    "print(f\"Accuracy: {accuracy:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The results aren't that good, which we can attribute to the small dataset size."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Parameter Efficient Fine-Tuning Techniques\n",
    "In this section, we demonstrate how to employ parameter-efficient fine-tuning (PEFT) techniques to adapt a pre-trained model for specific genomics tasks using the PEFT library.\n",
    "\n",
    "The LoraConfig object is instantiated to configure the PEFT parameters:\n",
    "\n",
    "+ task_type: Specifies the type of task, in this case, sequence classification (SEQ_CLS).\n",
    "+ r: The rank of the LoRA matrices.\n",
    "+ lora_alpha: Scaling factor for adaptive re-parameterization.\n",
    "+ target_modules: Modules within the model to apply PEFT re-parameterization (query, key, value in this example).\n",
    "+ lora_dropout: Dropout rate used during PEFT fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DNA_LM(\n",
       "  (model): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(5504, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0-11): 12 x BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear (768 -> 768)\n",
       "              (key): Linear (768 -> 768)\n",
       "              (value): Linear (768 -> 768)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear (768 -> 768)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear (768 -> 3072)\n",
       "            (intermediate_act_fn): GELU(approximate='none')\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear (3072 -> 768)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (classifier): Linear (768 -> 2)\n",
       ")"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Number of classes for your classification task\n",
    "num_labels = 2\n",
    "classification_model = DNA_LM(lm, num_labels)\n",
    "classification_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindnlp.peft import LoraConfig, TaskType\n",
    "\n",
    "peft_config = LoraConfig(\n",
    "    r=8,\n",
    "    lora_alpha=32,\n",
    "    target_modules=[\"query\", \"key\", \"value\"],\n",
    "    lora_dropout=0.01,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 442,368 || all params: 90,121,730 || trainable%: 0.49085608986866985\n"
     ]
    }
   ],
   "source": [
    "from mindnlp.peft import get_peft_model\n",
    "\n",
    "peft_model = get_peft_model(classification_model, peft_config)\n",
    "peft_model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModel(\n",
       "  (base_model): LoraModel(\n",
       "    (model): DNA_LM(\n",
       "      (model): BertModel(\n",
       "        (embeddings): BertEmbeddings(\n",
       "          (word_embeddings): Embedding(5504, 768, padding_idx=0)\n",
       "          (position_embeddings): Embedding(512, 768)\n",
       "          (token_type_embeddings): Embedding(2, 768)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (encoder): BertEncoder(\n",
       "          (layer): ModuleList(\n",
       "            (0-11): 12 x BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): lora.Linear(\n",
       "                    (base_layer): Linear (768 -> 768)\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.01, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear (768 -> 8)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear (8 -> 768)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (key): lora.Linear(\n",
       "                    (base_layer): Linear (768 -> 768)\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.01, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear (768 -> 8)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear (8 -> 768)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (value): lora.Linear(\n",
       "                    (base_layer): Linear (768 -> 768)\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.01, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear (768 -> 8)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear (8 -> 768)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                  )\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear (768 -> 768)\n",
       "                  (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear (768 -> 3072)\n",
       "                (intermediate_act_fn): GELU(approximate='none')\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear (3072 -> 768)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (classifier): Linear (768 -> 2)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "peft_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/60 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 1/60 [00:00<00:38,  1.54it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.865, 'learning_rate': 1.9666666666666666e-05, 'epoch': 0.08}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|▎         | 2/60 [00:01<00:34,  1.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8191, 'learning_rate': 1.9333333333333333e-05, 'epoch': 0.17}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▌         | 3/60 [00:01<00:33,  1.72it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7792, 'learning_rate': 1.9e-05, 'epoch': 0.25}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  7%|▋         | 4/60 [00:02<00:32,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6196, 'learning_rate': 1.866666666666667e-05, 'epoch': 0.33}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  8%|▊         | 5/60 [00:02<00:31,  1.75it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7699, 'learning_rate': 1.8333333333333333e-05, 'epoch': 0.42}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|█         | 6/60 [00:03<00:30,  1.79it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6927, 'learning_rate': 1.8e-05, 'epoch': 0.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 12%|█▏        | 7/60 [00:04<00:29,  1.80it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7386, 'learning_rate': 1.7666666666666668e-05, 'epoch': 0.58}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 13%|█▎        | 8/60 [00:04<00:29,  1.78it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8408, 'learning_rate': 1.7333333333333336e-05, 'epoch': 0.67}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 15%|█▌        | 9/60 [00:05<00:29,  1.75it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6235, 'learning_rate': 1.7e-05, 'epoch': 0.75}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 17%|█▋        | 10/60 [00:05<00:29,  1.72it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7194, 'learning_rate': 1.6666666666666667e-05, 'epoch': 0.83}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 18%|█▊        | 11/60 [00:06<00:29,  1.67it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6711, 'learning_rate': 1.6333333333333335e-05, 'epoch': 0.92}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 12/60 [00:07<00:28,  1.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8141, 'learning_rate': 1.6000000000000003e-05, 'epoch': 1.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                               \n",
      " 20%|██        | 12/60 [00:26<00:28,  1.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.74076908826828, 'eval_runtime': 3.9145, 'eval_samples_per_second': 3.066, 'eval_steps_per_second': 0.511, 'epoch': 1.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 22%|██▏       | 13/60 [00:27<05:11,  6.62s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8463, 'learning_rate': 1.5666666666666667e-05, 'epoch': 1.08}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 23%|██▎       | 14/60 [00:28<03:39,  4.78s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8359, 'learning_rate': 1.5333333333333334e-05, 'epoch': 1.17}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 25%|██▌       | 15/60 [00:28<02:37,  3.49s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7896, 'learning_rate': 1.5000000000000002e-05, 'epoch': 1.25}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 27%|██▋       | 16/60 [00:29<01:55,  2.62s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.5835, 'learning_rate': 1.4666666666666666e-05, 'epoch': 1.33}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 28%|██▊       | 17/60 [00:29<01:25,  2.00s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7319, 'learning_rate': 1.4333333333333334e-05, 'epoch': 1.42}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███       | 18/60 [00:30<01:05,  1.55s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.734, 'learning_rate': 1.4e-05, 'epoch': 1.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 32%|███▏      | 19/60 [00:30<00:51,  1.24s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7376, 'learning_rate': 1.3666666666666667e-05, 'epoch': 1.58}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 33%|███▎      | 20/60 [00:31<00:40,  1.02s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7403, 'learning_rate': 1.3333333333333333e-05, 'epoch': 1.67}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 35%|███▌      | 21/60 [00:31<00:34,  1.14it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6485, 'learning_rate': 1.3000000000000001e-05, 'epoch': 1.75}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 37%|███▋      | 22/60 [00:32<00:29,  1.28it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6661, 'learning_rate': 1.2666666666666667e-05, 'epoch': 1.83}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 38%|███▊      | 23/60 [00:32<00:25,  1.43it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6876, 'learning_rate': 1.2333333333333334e-05, 'epoch': 1.92}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 24/60 [00:33<00:23,  1.53it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7695, 'learning_rate': 1.2e-05, 'epoch': 2.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                               \n",
      " 40%|████      | 24/60 [00:51<00:23,  1.53it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.7271575927734375, 'eval_runtime': 4.0855, 'eval_samples_per_second': 2.937, 'eval_steps_per_second': 0.49, 'epoch': 2.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 42%|████▏     | 25/60 [00:51<03:29,  5.99s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8329, 'learning_rate': 1.1666666666666668e-05, 'epoch': 2.08}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 26/60 [00:52<02:27,  4.35s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7254, 'learning_rate': 1.1333333333333334e-05, 'epoch': 2.17}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 45%|████▌     | 27/60 [00:52<01:46,  3.22s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.73, 'learning_rate': 1.1000000000000001e-05, 'epoch': 2.25}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 47%|████▋     | 28/60 [00:53<01:17,  2.43s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6515, 'learning_rate': 1.0666666666666667e-05, 'epoch': 2.33}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 48%|████▊     | 29/60 [00:54<00:58,  1.87s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7709, 'learning_rate': 1.0333333333333335e-05, 'epoch': 2.42}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|█████     | 30/60 [00:54<00:43,  1.46s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.741, 'learning_rate': 1e-05, 'epoch': 2.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 52%|█████▏    | 31/60 [00:55<00:34,  1.18s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7195, 'learning_rate': 9.666666666666667e-06, 'epoch': 2.58}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 53%|█████▎    | 32/60 [00:55<00:27,  1.00it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7575, 'learning_rate': 9.333333333333334e-06, 'epoch': 2.67}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 55%|█████▌    | 33/60 [00:56<00:23,  1.16it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6374, 'learning_rate': 9e-06, 'epoch': 2.75}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 57%|█████▋    | 34/60 [00:56<00:19,  1.32it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6527, 'learning_rate': 8.666666666666668e-06, 'epoch': 2.83}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 58%|█████▊    | 35/60 [00:57<00:17,  1.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6743, 'learning_rate': 8.333333333333334e-06, 'epoch': 2.92}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 36/60 [00:57<00:16,  1.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7407, 'learning_rate': 8.000000000000001e-06, 'epoch': 3.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                               \n",
      " 60%|██████    | 36/60 [01:16<00:16,  1.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.7168386578559875, 'eval_runtime': 4.1825, 'eval_samples_per_second': 2.869, 'eval_steps_per_second': 0.478, 'epoch': 3.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 62%|██████▏   | 37/60 [01:17<02:25,  6.35s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.782, 'learning_rate': 7.666666666666667e-06, 'epoch': 3.08}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 63%|██████▎   | 38/60 [01:18<01:41,  4.63s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8069, 'learning_rate': 7.333333333333333e-06, 'epoch': 3.17}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 65%|██████▌   | 39/60 [01:18<01:12,  3.43s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7482, 'learning_rate': 7e-06, 'epoch': 3.25}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 67%|██████▋   | 40/60 [01:19<00:51,  2.59s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6373, 'learning_rate': 6.666666666666667e-06, 'epoch': 3.33}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 68%|██████▊   | 41/60 [01:20<00:37,  2.00s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7137, 'learning_rate': 6.333333333333333e-06, 'epoch': 3.42}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 70%|███████   | 42/60 [01:20<00:28,  1.58s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7217, 'learning_rate': 6e-06, 'epoch': 3.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 72%|███████▏  | 43/60 [01:21<00:21,  1.28s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6666, 'learning_rate': 5.666666666666667e-06, 'epoch': 3.58}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 73%|███████▎  | 44/60 [01:21<00:17,  1.07s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8221, 'learning_rate': 5.333333333333334e-06, 'epoch': 3.67}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 75%|███████▌  | 45/60 [01:22<00:14,  1.07it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6367, 'learning_rate': 5e-06, 'epoch': 3.75}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 77%|███████▋  | 46/60 [01:23<00:11,  1.20it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6161, 'learning_rate': 4.666666666666667e-06, 'epoch': 3.83}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 78%|███████▊  | 47/60 [01:23<00:09,  1.31it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6651, 'learning_rate': 4.333333333333334e-06, 'epoch': 3.92}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 48/60 [01:24<00:08,  1.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7697, 'learning_rate': 4.000000000000001e-06, 'epoch': 4.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                               \n",
      " 80%|████████  | 48/60 [01:42<00:08,  1.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.7099848389625549, 'eval_runtime': 3.9243, 'eval_samples_per_second': 3.058, 'eval_steps_per_second': 0.51, 'epoch': 4.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 82%|████████▏ | 49/60 [01:42<01:07,  6.10s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8004, 'learning_rate': 3.6666666666666666e-06, 'epoch': 4.08}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 83%|████████▎ | 50/60 [01:43<00:44,  4.43s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8062, 'learning_rate': 3.3333333333333333e-06, 'epoch': 4.17}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 85%|████████▌ | 51/60 [01:43<00:29,  3.25s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6965, 'learning_rate': 3e-06, 'epoch': 4.25}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 87%|████████▋ | 52/60 [01:44<00:19,  2.43s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6102, 'learning_rate': 2.666666666666667e-06, 'epoch': 4.33}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 88%|████████▊ | 53/60 [01:44<00:13,  1.87s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7549, 'learning_rate': 2.3333333333333336e-06, 'epoch': 4.42}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 90%|█████████ | 54/60 [01:45<00:08,  1.49s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7021, 'learning_rate': 2.0000000000000003e-06, 'epoch': 4.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 92%|█████████▏| 55/60 [01:46<00:06,  1.22s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.717, 'learning_rate': 1.6666666666666667e-06, 'epoch': 4.58}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 93%|█████████▎| 56/60 [01:46<00:04,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.7023, 'learning_rate': 1.3333333333333334e-06, 'epoch': 4.67}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 95%|█████████▌| 57/60 [01:47<00:02,  1.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6489, 'learning_rate': 1.0000000000000002e-06, 'epoch': 4.75}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 97%|█████████▋| 58/60 [01:47<00:01,  1.26it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6614, 'learning_rate': 6.666666666666667e-07, 'epoch': 4.83}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 98%|█████████▊| 59/60 [01:48<00:00,  1.40it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.6609, 'learning_rate': 3.3333333333333335e-07, 'epoch': 4.92}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60/60 [01:49<00:00,  1.50it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.8036, 'learning_rate': 0.0, 'epoch': 5.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                               \n",
      "100%|██████████| 60/60 [01:52<00:00,  1.88s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.7075120806694031, 'eval_runtime': 3.9299, 'eval_samples_per_second': 3.054, 'eval_steps_per_second': 0.509, 'epoch': 5.0}\n",
      "{'train_runtime': 112.9944, 'train_samples_per_second': 8.496, 'train_steps_per_second': 0.531, 'train_loss': 0.7251347482204438, 'epoch': 5.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=60, training_loss=0.7251347482204438, metrics={'train_runtime': 112.9944, 'train_samples_per_second': 8.496, 'train_steps_per_second': 0.531, 'train_loss': 0.7251347482204438, 'epoch': 5.0})"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Define training arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir='./results',\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    num_train_epochs=5,\n",
    "    weight_decay=0.01,\n",
    "    eval_steps=1,\n",
    "    logging_steps=1,\n",
    ")\n",
    "\n",
    "# Initialize Trainer\n",
    "trainer = Trainer(\n",
    "    model=peft_model.model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=valid_dataset,\n",
    "    tokenizer=tokenizer,\n",
    ")\n",
    "\n",
    "# Train the model\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Evaluate PEFT Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 12/12 [00:02<00:00,  4.08it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0 1 1 0 1 0 0 1 0 1 0 0 0 1 1 0 1 0 1 1 1 1 1 1 0 0 0 0 1 0 1 1 1 0 0 0 0\n",
      " 1 0 1 1 0 0 1 1 1 0 1 0 0 0 0 1 1 1 0 0 1 1 0 1 0 1 0 0 1 1 0 0 1 0 0 0 0\n",
      " 1 1 0 1 0 1 1 0 1 1 1 1 1 0 1 1 0 0 0 1 1 0 0 0 1 0 1 1 1 1 1 1 0 0 0 0 0\n",
      " 0 1 1 0 0 0 0 1 1 0 1 0 0 1 0 1 1 0 0 1 0 1 0 1 0 1 0 1 1 0 1 0 1 1 0 1 1\n",
      " 0 0 1 0 0 0 1 0 1 0 1 1 0 1 1 0 1 0 0 0 1 0 1 0 0 0 0 1 1 1 1 1 0 0 1 1 0\n",
      " 1 0 0 0 1 0 0]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Generate predictions\n",
    "\n",
    "predictions = trainer.predict(test_dataset)\n",
    "logits = predictions.predictions\n",
    "predicted_labels = logits.argmax(axis=-1)\n",
    "print(predicted_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.58\n"
     ]
    }
   ],
   "source": [
    "def calculate_accuracy(true_labels, predicted_labels):\n",
    "\n",
    "    assert len(true_labels) == len(predicted_labels), \"Arrays must have the same length\"\n",
    "    correct_predictions = np.sum(true_labels == predicted_labels)\n",
    "    accuracy = correct_predictions / len(true_labels)\n",
    "\n",
    "    return accuracy\n",
    "\n",
    "test_labels = []\n",
    "for data in test_dataset:\n",
    "    sequence, label = data\n",
    "    for single_label in label:\n",
    "        test_labels.append(single_label)\n",
    "\n",
    "accuracy = calculate_accuracy(test_labels, predicted_labels)\n",
    "print(f\"Accuracy: {accuracy:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the PEFT model achieves similar performance to the baseline model, demonstrating the effectiveness of PEFT in adapting pre-trained models to specific tasks with limited computational resources.\n",
    "\n",
    "With PEFT, we only train 442,368 parameters, which is 0.49% of the total parameters in the model. This is a significant reduction in computational resources compared to training the entire model from scratch.\n",
    "\n",
    "We can improve the results by using a larger dataset, fine-tuning the model for more epochs or changing the hyperparameters (rank, learning rate, etc.)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "neoming_vivit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
